#!/usr/bin/env python
# encoding: utf-8
"""
Generator.py

Created by Alakazam on 2008-05-26.
Copyright (c) 2008. All rights reserved.
"""

from math import cos, sin, sqrt, radians
from numpy import *
from Numeric import array

from pca_module import *

from RandomArray import *

import psyco
psyco.full()

class Generator(object):
  """This class generates sample multidimensional data"""
  
  # List of existing generators
  generators = ["Gaussian", "SwissRoll", "4DSwissRoll", "SwissHole",
                "CornerPlanes", "ToroidalHelix", "3DClusters",
                "FromImages"]
  
  def __init__(self, generator="SwissRoll", n=100000, param=1, noise=1, dim=3):
    # Generator to use
    if generator not in Generator.generators:
      raise NotImplementedError, 'This generator has not been implemented'
    self.generator = generator
    
    # Number of points
    self.n         = n
    # Extra parameter to provide to the generator
    self.param     = param
    # Noise to add
    self.noise     = noise
    # Desired dimensionality of the data set
    self.dim       = dim
  
  def generateData(self):
    # The generator's return value is used to generate the color
    # reprensentation if != None, otherwise it is the generator's
    # responsibility to correctly initialise self.colors
    index = self.__getattribute__('generate' + self.generator)()
    
    if index != None:
      # Generate the colors from a one dimensional representation of the data
      # points
      self.colors = array([[pos, 1.0 - pos, 0.0] for pos in index])
    
    currentDim = len(self.data[0])
    
    if currentDim < self.dim:
      # The initial data set does not have enough dimensions
      self.data = concatenate((self.data, zeros((self.n, self.dim - currentDim))), 1)
    
    if self.noise > 0:
      self.data += float(self.noise) * random((self.n, self.dim))
  
  def generateGaussian(self):
    """Gaussian randomly sampled data generator"""
    r = random((self.n, 2)) - 0.5
    
    self.data = concatenate((
      r,
      reshape((1 / ((1.0/self.param) **2 * 2 * pi)) * exp(-(r[:, 0] **2 + r[:, 1] **2) / (2 * (1.0/self.param) **2)), (-1, 1))),
      1)
    
    # Generate indexes between 0 and 1 for colorization
    index = self.data[:, 2]
    m     = min(index)
    
    return (index - m) / (max(index) - m)
  
  def generateSwissRoll(self):
    """Swiss roll generator"""
    phi = random(self.n)
    t = (3 * pi / 2) * (1 + 2 * phi);
    height = 21 * random(self.n);
    self.data = transpose(array([t * cos(t), height, self.param * t * sin(t)]))
    
    return phi
  
  def generate4DSwissRoll(self):
    """4D Swiss roll generator"""
    phi = random(self.n)
    t = (3 * pi / 2) * (1 + 2 * phi)
    height = 21 * random(self.n)
    depth  = 21 * random(self.n)
    self.data = transpose(array([t * cos(t), height, self.param * t * sin(t), depth]))
    
    return phi
  
  def generateSwissHole(self):
    """Swiss Roll w/ hole example taken from Donoho & Grimes"""
    r = random(2 * self.n)
    t = (3 * pi / 2) * (1 + 2 * r)
    
    height = 21 * random(2 * self.n)
    
    data = reshape(
      [[s * cos(s), h, self.param * s * sin(s), r]
        for (s, h, r) in zip(t, height, r)
        if s < 9 or s > 12 or h < 9 or h > 14],
      (-1, 4))[0:self.n]
    
    self.data = data[:, 0:3]
    
    return data[:, 3]
  
  def generateCornerPlanes(self):
    """Corner Planes generator"""
    
    xMax        = int(floor(sqrt(self.n)));
    yMax        = int(ceil(self.n / xMax)) + 1;
    cornerPoint = int(floor(yMax / 2));
    
    data = reshape(
      [[x,
        y if y <= cornerPoint
        else cornerPoint + (y - cornerPoint) * cos(pi * self.param / 180),
        0 if y <= cornerPoint
        else (y - cornerPoint) * sin(pi * self.param / 180),
        float(y) / yMax]
        for x in range(0, xMax) for y in range(0, yMax)],
      (-1, 4))[0:self.n]
    
    self.data = data[:, 0:3]
    
    return data[:, 3]
  
  def generateToroidalHelix(self):
    """Toroidal Helix generator by Coifman & Lafon"""
    
    r = array([float(i) / self.n * 2 * pi for i in range(1, self.n + 1)])
    
    self.data = transpose([(2 + cos(8 * self.param * r)) * cos(r), (2 + cos(8 * self.param * r)) * sin(r), sin(8 * r)])
    
    return [float(i) / self.n for i in range(0, self.n)]
  
  def generate3DClusters(self):
    """3D Clusters generator"""
    nClusters = max(2, int(self.param))
    centers   = 10 * random((nClusters, 3))
    
    minDist   = float('Inf')
    for c1 in centers:
      for c2 in centers:
        if c1 == c2:
          continue
        
        d = (c1 - c2) **2
        if d < minDist:
          minDist = d
    
    # Cluster radius
    r = sqrt(minDist) / 3
    
    # Number of points not in clusters or connecting lines
    nPts = self.n - nClusters - (nClusters - 1) * 9
    
    # Number of points per cluster (rounded up)
    nPtsPerCluster = int(float(nPts) / nClusters) + 1
    
    (self.data, index) = (centers, zeros(nClusters))
    
    for i in range(0, nClusters):
      self.data = concatenate(
        (self.data,
         reshape([centers[i] + r * (random((1, 3)) - 0.5) for j in range(0, nPtsPerCluster)], (-1, 3))))
      
      index = concatenate(
        (index,
         zeros(nPtsPerCluster) + float(i) / nClusters))
    
    # Crop cluster points to nPts :
    self.data = self.data[0:(nPts + nClusters)]
    index = index[0:(nPts + nClusters)]
    
    # Generate connecting segments :
    ts = [t / 10. for t in range(1, 10)]
    
    lines = concatenate([[centers[i] * t + centers[i + 1] * (1. - t) for t in ts] for i in range(0, nClusters - 1)])
    
    self.data = concatenate((self.data, lines))
    
    index = concatenate(
      (index,
       zeros(9 * (nClusters - 1))))
    
    return index
  
  def generateFromImages(self):
    from PIL import Image, ImageFilter, ImageDraw, ImageChops
    import os
    from colorsys import hsv_to_rgb
    
    self.data = array([Image.open(imageFile).convert("L").getdata() for imageFile in self.param])
    
    nbImages = self.data.shape[0]
    
    # Group images by 10
    self.colors = array([list(hsv_to_rgb((i / 10) / (nbImages / 10.), 0.5, 0.8)) for i in range(0, nbImages)])
    
    return None
  
  def getColors(self):
    return self.colors
  
  def displayData(self, k):
    """Display the colored data generated by this generator
    
    k is the number of nearest neighbours to display
    """
    
    from OpenGLDataView import OpenGLDataView
    
    OpenGLDataView(self.data, knn=k, colors=self.colors).run()
